"""
WebSocket JWT认证中间件
用于处理WebSocket连接的JWT Token认证
"""
import jwt
from django.conf import settings
from django.contrib.auth import get_user_model
from django.contrib.auth.models import AnonymousUser
from channels.middleware import BaseMiddleware
from channels.db import database_sync_to_async
from urllib.parse import parse_qs
import logging

logger = logging.getLogger(__name__)
User = get_user_model()


@database_sync_to_async
def get_user_by_token(token):
    """通过JWT token获取用户"""
    try:
        # 使用rest_framework_simplejwt来解码token
        from rest_framework_simplejwt.tokens import AccessToken
        
        # 验证和解码token
        access_token = AccessToken(token)
        user_id = access_token['user_id']
        
        # 查询用户
        user = User.objects.get(id=user_id)
        return user
        
    except Exception as e:
        logger.warning(f"JWT token认证失败: {e}")
        return AnonymousUser()


class JWTAuthMiddleware(BaseMiddleware):
    """
    JWT认证中间件
    从WebSocket连接的查询参数中提取token并进行认证
    """
    
    async def __call__(self, scope, receive, send):
        # 只处理WebSocket连接
        if scope["type"] == "websocket":
            # 从查询参数中获取token
            query_params = parse_qs(scope["query_string"].decode())
            token = query_params.get("token")
            
            if token and len(token) > 0:
                token_value = token[0]  # 获取第一个token值
                logger.info(f"WebSocket连接包含token: {token_value[:20]}...")
                
                # 通过token获取用户
                user = await get_user_by_token(token_value)
                scope["user"] = user
                
                if user.is_authenticated:
                    logger.info(f"WebSocket JWT认证成功: {user.username}")
                else:
                    logger.warning("WebSocket JWT认证失败: 用户未认证")
            else:
                logger.warning("WebSocket连接缺少token参数")
                scope["user"] = AnonymousUser()
        
        return await super().__call__(scope, receive, send)


def JWTAuthMiddlewareStack(inner):
    """
    JWT认证中间件栈
    用于替换默认的AuthMiddlewareStack
    """
    return JWTAuthMiddleware(inner)